#!/usr/bin/env python
# Script that contains a function to import MOS forecast data from the IA state archive
# and write it to a class for mos forecasts
from datetime import datetime, timedelta
import sys
sys.path.append('..')
#infile = 'KPVD_GFS_MOS.csv'
from INFILE import siteid, maindir,startdate


# first, a class for the MOS values we want
class mosval:
    def __init__(self,valid,max,min,wind,precip):
        self.valid_date = valid
        self.max_temp = max
        self.min_temp = min
        self.max_wind = wind
        self.precip = precip

        
def process_mos_archive(infile):
    # Function to take the infile and read in values
    mos_dict = {}
    # Open file. First line in the file contains all info, ignore that.
    file = open(infile,'r').readlines()[1:]
    
    linenum = 1
    for line in file:
        try: 
            station,model,runtime,ftime,x_n,tmp,dpt,cld,wdr,wsp,p06,p12,q06,q12,t06,t12,snw,cig,vis,obv,poz,pos,typ = line.split(',')[:-1]    
        except:
            print line
            station,model,runtime,ftime,x_n,tmp,dpt,cld,wdr,wsp,p06,p12,q06,q12,t06,t12,snw,cig,vis,obv,poz,pos,typ = line.split(',')[:-1]    
               
         
        # All right.  We have the data.  Use datetime to extrac the times.
        inittime = datetime.strptime(runtime,'%Y-%m-%d %H:00:00+00')
        fcsttime = datetime.strptime(ftime,'%Y-%m-%d %H:00:00+00')
        
        # Quick correction for first date
        if linenum == 1:
            curdate = fcsttime
            curinittime = inittime
            wind_list = []
            precip_list = []
            high = None
            low = None
        
        # Now we check to see if it's a new day
        if fcsttime.strftime('%Y%m%d') != curdate.strftime('%Y%m%d') and int(fcsttime.strftime('%H')) == 6:
            # Since we go from approximately 6Z to 6Z, reset at hour 6
            # Store the previous day of data
            #print "Curdate", curdate.strftime('%Y-%m-%d')
            #print "Curinittime",curinittime.strftime('%Y-%m-%d')
            #print "Fcsttime",fcsttime.strftime('%Y-%m-%d')
            fcst_lead = curdate-curinittime
            if curinittime in mos_dict.keys():
                pass
            else:
                mos_dict[curinittime] = {}

            mos_dict[curinittime][fcst_lead.days] = mosval(curdate,max=high,min=low,wind=max(wind_list),precip=sum(precip_list))
 
            # Check
            #if curdate.strftime('%Y%m%d') == '20120819':
            #    print "FOUND!"
            #    print "INIT:", curinittime.strftime('%Y-%m-%d %HZ'), "LEAD DAYS:", fcst_lead.days
            #    print mos_dict[curinittime][fcst_lead.days].valid_date.strftime('%Y-%m-%d')
            #    print mos_dict[curinittime][fcst_lead.days].max_temp
            #    print mos_dict[curinittime][fcst_lead.days].min_temp
            #    print mos_dict[curinittime][fcst_lead.days].max_wind
            #    print "KEYS BEFORE (%s):" % infile
            #    print mos_dict[datetime(2012,8,19,12)].keys()
            #    raw_input()

            curdate = fcsttime
            curinittime = inittime
            wind_list = []
            precip_list = []
            high = None
            low = None
        
        # Now just grab the numbers
        if x_n != '':
            #print "Finding data!"
            if fcsttime.strftime('%H') == '12':
                low = int(x_n)
            else:
                high = int(x_n)

        wind_list.append(int(wsp))
        if q06 != '':
            precip_list.append(int(q06))
        linenum = linenum + 1   
    outfile = open(infile[:-4] + '.pickle','w')
    import cPickle
    cPickle.dump(mos_dict,outfile)
    outfile.close()    

def read_mos_dict(dictfile):
    # Function to read certain values
    import cPickle
    from datetime import datetime
    mos_dict = cPickle.load(open(dictfile,'r'))

    fcst_highs = []
    fcst_lows = []
    fcst_winds = []
    fcst_precip = []
    dates = []
    # Now the sorting begins  
    dkeys = mos_dict.keys()
    dkeys.sort()
    outdict = {}
    for day in dkeys:
        if day.strftime('%H') == '12':
            # Only want 12Z runs
            try:
                high = mos_dict[day][0].max_temp
            except:
                print day, dictfile
                print mos_dict[day].keys()
                high = mos_dict[day][0].max_temp
     
            low = mos_dict[day][0].min_temp
            wind = mos_dict[day][0].max_wind
            precip = mos_dict[day][0].precip
            datval = mos_dict[day][0].valid_date

            outdict[datval] = {}
            outdict[datval]['high'] = high
            outdict[datval]['low'] = low
            outdict[datval]['wind'] = wind
            outdict[datval]['precip'] = precip

    #print dates
    #print fcst_highs
    return outdict



def update_mos_archive(siteid,model,datein=None):
    # Uses the handy web address interface of the IA state database to update
    # the mos archive
    from datetime import datetime, timedelta
    from urllib import urlopen
    run_interval = timedelta(hours=12)   # MOS updates every 12 hours

    # First, see what run we left off at
    try:
        archive_file = open('%s/site_data/%s_%s_MOS.csv' % (maindir,siteid.upper(),model.upper()),'ra')
        last_runtime_text = archive_file.readlines()[-1].split(',')[2]
        # Convert this to a datetime
        last_runtime = datetime.strptime(last_runtime_text,'%Y-%m-%d %H:00:00+00')
        archive_file.close() 
    except:
        # We didn't find a mos archive.  Try downloading one beginning at the startdate
        datein = startdate


    def retrieve_file(fmt_curtime):
        # Get the file for the curtime
        webfile = urlopen('http://mesonet.agron.iastate.edu/mos/csv.php?station=%s&runtime=%s&model=%s' % (siteid,fmt_curtime,model))
        infile = webfile.readlines()
        webfile.close()
        if len(infile) <2:
            # The file download was not successful
            print "WARNING: Data missing for %s MOS at time" % (model.upper()), fmt_curtime
        with open('%s/site_data/%s_%s_MOS.csv' % (maindir,siteid.upper(),model.upper()), 'a') as archive_file:
            for line in infile[1:]:
                archive_file.write(line)

    print "Updating %s MOS archive for %s" % (model.upper(), siteid.upper())
    if datein != None:
        fmt_curtime = datein.strftime('%Y-%m-%d%%20%H:00')
        retrieve_file(fmt_curtime)

    else:
        # Now add times until today
        now = datetime.utcnow()
        curtime = last_runtime + run_interval
        while curtime < now:
            fmt_curtime = curtime.strftime('%Y-%m-%d%%20%H:00')
            retrieve_file(fmt_curtime)
            curtime = curtime + run_interval


    # Finally, reprocess the text file into the dictionary
    process_mos_archive('%s/site_data/%s_%s_MOS.csv' % (maindir,siteid.upper(),model.upper()))
    #print last_runtime
 







update_mos_archive(siteid.upper(),'GFS') #,datetime.strptime(indate,'%Y%m%d%H'))
print "GFS Done"
update_mos_archive(siteid.upper(),'NAM')
print "NAM Done"
#dictfile = 'pdict_' + infile
            
#process_mos_archive(infile)
#read_mos_dict(dictfile)

